#####################################
##       PEinjector/regwrite       ##
#####################################

import winreg
import log
import os


def add_path(path: str) -> None:  # 写注册表添加path
    log.info(f"add path \"{path}\"")
    with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
                        r"SYSTEM\CurrentControlSet\Control\Session Manager\Environment", 0, winreg.KEY_READ) as key:
        old_path = winreg.QueryValueEx(key, "Path")[0]
    if old_path[-1] != ';':
        old_path = old_path+";"
    new_path = old_path+path.replace("/", "\\")
    with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
                        r"SYSTEM\CurrentControlSet\Control\Session Manager\Environment", 0, winreg.KEY_WRITE) as key:
        winreg.SetValueEx(key, "Path", 0, winreg.REG_EXPAND_SZ, new_path)


def add_reg(regpath: str) -> None:  # 导入注册表文件
    log.info(f"add reg \"{regpath}\"")
    assert os.system(f"reg import \"{regpath.replace("/", "\\")}\">nul") == 0


def refresh_path() -> None:  # Thank for bbs.bathome.net/thread-62029-1-1.html and dllcall.exe
    log.info("refresh path env")
    os.system(
        f"start /I /B {os.path.dirname(__file__).replace("src", "tool")}\\dllcalls.bat>nul")
